package com.rockcypress.common.agent.utils;

import com.jcraft.jsch.ConfigRepository;
import com.jcraft.jsch.JSch;
import com.jcraft.jsch.JSchException;
import com.jcraft.jsch.Session;
import com.rockcypress.common.agent.dto.PortForwardDTO;

import java.io.*;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;

/**
 * @author bolin
 * @create 2020-05-12 5:05 下午
 * @desc
 */
public class JSchUtils {

    final private static JSch jsch = new JSch();

    private volatile static String currConfigFile = null;

    public final static String DEFAULT_CONFIG_FILE = "~/.ssh/config";

    public final static String DEFAULT_CONFIG_DATA = "forwardMap.data";

    private final static String DEFAULT_PRIVATE_KEY_FILE = "~/.ssh/id_rsa";

    private final static String IDENTITY_FILE_KEY = "identityfile";

    public final static Properties properties = new Properties();

    static {
        properties.put("StrictHostKeyChecking", "no"); // 不校验第一次链接的公钥询问，与ip变化或其他场景。 注意：此举会导致中间人攻击可行，不过这个本来就是私有操作，所以放开
    }

    private final static Map<String, Session> sshSessionMap = new ConcurrentHashMap<>();

    private final static Map<String, Set<PortForwardDTO>> forwardMap = new HashMap<>();


    synchronized public static String forward(String host, Integer localPort, String remoteHost, Integer remotePort) throws JSchException {
        if (!sshSessionMap.containsKey(host)) {
            throw new IllegalArgumentException("[" + host + "]通道未打开无法进行转发, 可以先运行: open " + host);
        }
        ConfigRepository.Config sshConfig = jsch.getConfigRepository().getConfig(host);
        sshSessionMap.get(host).setPortForwardingL(localPort, remoteHost, remotePort);
        addForward(host, localPort, remoteHost, remotePort);
        return "正在监听本地端口[" + localPort + "],数据通过[" + host + "]转发至[" + remoteHost + ":" + remotePort + "]";
    }

    public static String connectionByPK(String host) throws JSchException {
        if (sshSessionMap.containsKey(host)) {
            throw new IllegalArgumentException("通道已打开无需再次链接[" + host + "] ,可以使用forward进行转发");
        }
        ConfigRepository.Config sshConfig = jsch.getConfigRepository().getConfig(host);
        if (sshConfig.getValue(IDENTITY_FILE_KEY) == null || sshConfig.getValue(IDENTITY_FILE_KEY).length() == 0) {
            if (defaultPrivateKeyFileExists()) {
                jsch.addIdentity(DEFAULT_PRIVATE_KEY_FILE);
            } else {
                throw new IllegalArgumentException(getCurrConfigFile() + " 未配置[" + host + "]的[" + IDENTITY_FILE_KEY + "]属性");
            }
        } else {
            jsch.addIdentity(sshConfig.getValue(IDENTITY_FILE_KEY));
        }

        Session session = jsch.getSession(sshConfig.getUser(), sshConfig.getHostname(), sshConfig.getPort());
        session.setConfig(properties);
        session.connect();
        sshSessionMap.put(host, session);
        return "[" + host + "]ssh通道已打开, 可以通过forward命令进行转发";
    }

    public static void connectionByPW(String host, String password) throws JSchException {
        if (sshSessionMap.containsKey(host)) {
            throw new IllegalArgumentException("通道已打开无需再次链接: " + host);
        }
        ConfigRepository.Config sshConfig = jsch.getConfigRepository().getConfig(host);
        Session session = JSchUtils.getInstance().getSession(sshConfig.getUser(), sshConfig.getHostname(), sshConfig.getPort());
        session.setPassword(password);
        session.connect();
        sshSessionMap.put(host, session);
    }

    public static JSch getInstance() {
        return jsch;
    }

    public static Object close(String host, Integer localPort) throws JSchException {
        if (!sshSessionMap.containsKey(host)) {
            throw new IllegalArgumentException("[" + host + "]: 通道未打开无需关闭");
        }
        Session session = sshSessionMap.get(host);
        if (localPort != null) {
            PortForwardDTO dto = null;
            Iterator<PortForwardDTO> iterator = forwardMap.get(host).iterator();
            while (iterator.hasNext()) {
                dto = iterator.next();
                if (dto.getLocalPort().equals(localPort)) {
                    session.delPortForwardingL(localPort);
                    return forwardMap.get(host).remove(dto);
                }
            }
            throw new IllegalArgumentException("端口未打开无需关闭: [" + localPort + "]");
        } else {
            List<PortForwardDTO> list = getForwardList(host);
            session.disconnect();
            sshSessionMap.remove(host); // 移除 session
            forwardMap.remove(host);// 移除转发记录
            return list;
        }
    }

    synchronized private static void addForward(String host, int localPort, String remoteHost, int remotePort) {
        Set<PortForwardDTO> set = forwardMap.get(host);
        if (set == null) {
            set = new HashSet<>();
            forwardMap.put(host, set);
        }
        set.add(new PortForwardDTO(host, localPort, remoteHost, remotePort));
    }

    public static List<PortForwardDTO> getForwardList(String host) {
        List<PortForwardDTO> list = new ArrayList<>();
        if (!forwardMap.containsKey(host)) {
            return list;
        }
        return forwardMap.get(host).stream().collect(Collectors.toList());
    }

    public static List<String> getSessionList() {
        return sshSessionMap.keySet().stream().collect(Collectors.toList());
    }

    public static List<PortForwardDTO> save(String host) throws IOException, ClassNotFoundException {
        List<PortForwardDTO> list = getForwardList(host);
        if (list.isEmpty()) {
            return list;
        }
        Map<String, List<PortForwardDTO>> map = reload();
        map.put(host, list);
        File saveFile = new File(DEFAULT_CONFIG_DATA);
        if (!saveFile.exists()) {
            saveFile.createNewFile();
        }
        ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(saveFile));
        oos.writeObject(map);
        oos.flush();
        oos.close();
        return list;
    }

    public static Map<String, List<PortForwardDTO>> reload() throws IOException, ClassNotFoundException {
        Map<String, List<PortForwardDTO>> map = new HashMap<>();
        File saveFile = new File(DEFAULT_CONFIG_DATA);
        if (!saveFile.exists()) {
            return map;
        }
        ObjectInputStream ois = new ObjectInputStream(new FileInputStream(saveFile));
        Map<String, List<PortForwardDTO>> fileMap = (Map<String, List<PortForwardDTO>>) ois.readObject();
        map.putAll(fileMap);
        ois.close();
        return map;
    }

    public static Object load(String host) throws IOException, ClassNotFoundException, JSchException {
        List<PortForwardDTO> list = reload().get(host);
        if (!sshSessionMap.containsKey(host)) {
            connectionByPK(host);
        }
        for (PortForwardDTO dto : list) {
            forward(dto.getHost(), dto.getLocalPort(), dto.getRemoteHost(), dto.getRemotePort());
        }
        return list;
    }

    public static String getCurrConfigFile() {
        return currConfigFile == null ? DEFAULT_CONFIG_FILE : currConfigFile;
    }

    public static void setCurrConfigFile(String currConfigFile) {
        JSchUtils.currConfigFile = currConfigFile;
    }

    private static boolean defaultPrivateKeyFileExists() {
        return new File(DEFAULT_PRIVATE_KEY_FILE).exists();
    }
}
